# encoding: utf-8
# $Rev: 334 $

"""
Un repositorio de descarga de archivos para el nodo P2P

El repositorio permite realizar operaciones sobre un conjunto de archivos a
los cuales se refiere por el MD5 del contenido. De este modo puede almacenar
diferentes archivos con el mismo nombre, siempre y cuando tengan diferente
contenido.

Además permite almacenar archivos "parciales", es decir archivos de los
cuales se tiene solo un subconjunto de los fragmentos totales. A los archivos no
parciales les llamaremos "completos".

Para inicializar un repositorio hay que construir un FileDatabase, dando
como argumento el directorio donde se van a almacenar los archivos. Para que
la base funcione, el directorio debe existir, y contener un subdirectorio
llamado 'partial'. No deberían introducirse archivos en este directorio fuera
de los que son creados por las operaciones del TAD FileDatabase.

>>> import os
>>> os.system('rm -rf download; mkdir download; mkdir download/partial')
0
>>> repositorio = FileDatabase('download')
>>> repositorio
file_db.FileDatabase('download')

Se pueden meter al repositorio archivos externos usando la operación 
import_file:

>>> open('ejemplo.txt', 'w').write('Este es el contenido de un archivo de ejemplo')
>>> repositorio.import_file('ejemplo.txt')

La operación list permite obtener una lista de los md5s de los archivos
*completos* contenidos:

>>> repositorio.list()
['2c43ed66f6869293c54e9e33b1a31cb0']

Usando los md5 de los archivos contenidos, se pueden obtener objetos de tipo
File(), que representan cada uno de los archivos (ya sean completos o parciales)
en la base. Para obtener un archivo, se usa la operación get()

>>> repositorio.get('2c43ed66f6869293c54e9e33b1a31cb0')
<file_db.File object; md5=2c43ed66f6869293c54e9e33b1a31cb0 path=ejemplo.txt complete>

La operacion get() devuelve None si se pide un md5 inexistente
>>> print repositorio.get('12345612345612345612345612345678')
None

Otra forma de crear un objeto File dentro del repositorio es usando la
operación new, tipicamente antes de empezar a descargar un archivo. Si bien
no hace falta el contenido para crear un file, es necesario ya tener algo de
información básica: md5, nombre, y tamaño total en bytes:

>>> descarga = repositorio.new('ac3461059ffc2095e86fd80e8b5fc280', 'vmlinuz-2.6.31-21-generic', 3901472)
>>> descarga
<file_db.File object; md5=ac3461059ffc2095e86fd80e8b5fc280 path=vmlinuz-2.6.31-21-generic partial>

El archivo recién creado se marca como "parcial". Por lo tanto sigue sin
aparecer en el resultado de list():

>>> repositorio.list()
['2c43ed66f6869293c54e9e33b1a31cb0']

Existe un metodo mark_as_complete que transforma archivos parciales en
completos, pero solo si el md5 coincide. Si el md5 no coincide devuelve None:

>>> print repositorio.mark_as_complete('ac3461059ffc2095e86fd80e8b5fc280')
None
>>> descarga.is_partial()
True
>>> descarga2 = repositorio.new('57bb4b485f52628aab93cbde76d322d0', 'prueba.txt', 9)
>>> descarga2.set_index_info(['57bb4b485f52628aab93cbde76d322d0'])
>>> descarga2.put_slice(1, 'contenido') # Eso es contenido con el tamaño y md5 de la línea anterior
>>> repositorio.mark_as_complete('57bb4b485f52628aab93cbde76d322d0')
<file_db.File object; md5=57bb4b485f52628aab93cbde76d322d0 path=prueba.txt complete>
>>> descarga2.is_partial()
False

Cada objeto archivo tiene algunos accesores para obtener su información:

>>> descarga = repositorio.get('ac3461059ffc2095e86fd80e8b5fc280')
>>> descarga.get_metadata()
('vmlinuz-2.6.31-21-generic', 3901472L)
>>> descarga.get_fragment_count()
30L
>>> descarga.is_partial()
True

"""



import constants
import os.path, os
import glob
import hashlib
import shutil
import logging

__all__ = ['FileDatabase']

def _find_file (path, md5):
    pattern = os.path.join (path, md5+".*")
    matches = glob.glob (pattern)
    if matches:
        assert len(matches)==1
        return matches[0]

class FileDatabase(object):
    """
    Repositorio de archivos descargados
    """
    def __init__ (self, path):
        """
        Inicializa un repositorio asociado al path dado. El path no
        necesariamente existe al momento de inicialización pero debe existir
        al momento de realizar cualquier operación.
        
        En el directorio, la implementación actual almacena
        
        path/             # Ruta base del repositorio
            md5.nombre    # Cada uno de los archivos completos contenidos, se agrega el md5 como prefijo
            md5.nombre
            partial/
               md5.nombre # Cada uno de los archivos parciales contenidos, se agrega el md5 como prefijo
               md5.nombre # Cada uno de los archivos parciales contenidos, se agrega el md5 como prefijo              
        """
        self.path = path # folder with complete files
        self.partial_path = os.path.join (path, 'partial')
        self.cache = {} # internal cache. Maps str(md5 hashcode) -> File()

    def __repr__ (self):
        return 'file_db.FileDatabase(%s)' % repr(self.path)

    def get (self, file_md5):
        """Obtiene el objeto File() asociado al md5 dado, o None"""
        if file_md5 in self.cache:
            return self.cache[file_md5]
        # Buscar archivo completo
        fn = _find_file (self.path, file_md5)
        if fn is not None:
            self.cache[file_md5] = File(fn)
            return self.cache[file_md5]
        # Buscar en descargas parciales
        fn = _find_file (self.partial_path, file_md5)
        if fn is not None:
            self.cache[file_md5] = File(fn, partial=True)
            return self.cache[file_md5] # file instance

    def list(self):
        """Lista de MD5s de todos los archivos completos en el repositorio"""
        result = os.listdir (self.path)
        return [fn[:32] for fn in result if fn != 'partial' and not fn.startswith('.')]

    def new (self, md5, name, size):
        """
        Crea un nuevo objeto File(), parcial, sin datos, cuyos metadatos son
        los argumentos dados
        """
        fname = os.path.join (self.partial_path, "%s.%s" % (md5, name))
        f = open (fname, 'ab')
        f.truncate (size)
        return self.get (md5)

    def mark_as_complete(self, md5):
        """
        Marca archivos descargados como completos.
        
        Chequea el md5. Si no coincide con lo que debería ser, devuelve None. Si
        coincide, devuelve el objeto File() actualizado
        """
        assert self.get(md5) is not None
        file = self.get(md5)
        assert file.is_partial()
        # Check md5
        f = open(file._path, 'rb')
        h = hashlib.md5()
        buf = f.read(4096)
        while buf:
            h.update(buf)
            buf = f.read(4096)
        f.close()
        if h.hexdigest() != md5:
            # The file is not complete, or corrupt; md5s do not match
            return None
        fname = os.path.join (self.path, os.path.basename(file._path))
        os.rename(file._path, fname)
        file._make_non_partial(fname)
        return file
        
    def import_file (self, path):
        """Importa un archivo externo, completo, al repositorio"""
        h = hashlib.md5()
        f = open(path, 'rb')
        buf = f.read(4096)
        while buf:
            h.update(buf)
            buf = f.read(4096)
        f.close()
        target = h.hexdigest()+ "." + os.path.basename (path)
        ptarget = os.path.join (self.partial_path, target)
        shutil.copyfile (path, ptarget)
        target = os.path.join (self.path, target)
        shutil.move (ptarget, target)
        self.cache[h.hexdigest()] = File (target)

class File (object):
    """
    Representación de cada uno de los archivos dentro del repositorio
    """
    def __init__ (self, path, partial=False):
        """
        Constructor. Nunca debería llamarse directamente; en vez deben generarse
        nuevos files usando las operaciones new() y get() de FileDatabase
        """
        self._path = path
        st = os.stat (path)
        self._size = st.st_size
        self.md5 = os.path.basename(path)[:32]
        self._name = os.path.basename(path)[32+1:]
        self._partial = partial
        self._index_fragments = None

    def __repr__(self):
        if self._partial:
            label ='partial'
        else:
            label='complete'
        return '<file_db.File object; md5=%s path=%s %s>' % (self.md5, self._name, label)

    def get_metadata(self):
        """
        Metadatos (nombre, tamaño) del archivo representado
        """
        return self._name, self._size

    def get_fragments (self):
        """
        Lista de los MD5 de cada uno de los fragmentos que el archivo tiene en
        guardados actualmente. Para una archivo parcial, estos MD5s pueden
        diferir de los MD5s que debería tener al archivo al completarse
        """
        f = open (self._path, 'rb')
        h = hashlib.md5()
        result = []

        buf = f.read(constants.FRAGMENT_SIZE)
        while buf:
            h.update(buf)
            result.append (h.hexdigest())
            h = hashlib.md5()
            buf = f.read(constants.FRAGMENT_SIZE)
        return result

    def get_slice (self, index):
        """
        Obtiene el fragmento con el indice dado del disco. Para archivos
        parciales, notese que el contenido puede no ser el correcto.
        """
        assert 0 < index <= self.get_fragment_count()
        f = open (self._path, 'rb')
        f.seek ((index-1)*constants.FRAGMENT_SIZE)
        return f.read(constants.FRAGMENT_SIZE)

    def get_fragment_count(self):
        """Cantidad de fragmentos que tiene este archivo"""
        return (self._size + (constants.FRAGMENT_SIZE - 1)) // constants.FRAGMENT_SIZE

    def is_partial(self):
        """True cuando el archivo es parcial, False cuando es completo"""
        return self._partial

    def set_index_info (self, fragments):
        """
        Este método permite asociar al archivo la información sobre MD5s de
        fragmentos que debería tener (usualmente obtenidos del índice). Con
        esta información, se puede verificar si un fragmento esta correctamente
        descargado o no.
        """
        assert len(fragments) == self.get_fragment_count()
        self._index_fragments = fragments

    def has_index_info (self):
        """True luego de que se ha llamado a set_index_info"""
        return not self._partial or self._index_fragments is not None

    def get_index_info (self):
        """Obtiene la información sobre MD5s seteada con set_index_info"""
        assert self.has_index_info()
        return self._index_fragments

    def put_slice (self, index, data):
        """Almacena el fragmento data en la posición index del archivo"""
        assert 0 < index <= self.get_fragment_count()
        assert len(data) <= constants.FRAGMENT_SIZE # Los datos no superan el tamaño de fragmento
        assert len(data) == constants.FRAGMENT_SIZE or index == self.get_fragment_count() # Solo el ultimo fragmento puede ser mas corto
        # Check data
        h = hashlib.md5()
        h.update (data)
        assert h.hexdigest() == self._index_fragments[index-1]
        # Store!
        f = open (self._path, 'r+b')
        f.seek ((index-1)*constants.FRAGMENT_SIZE)
        f.write (data)
        f.close()
        logging.debug("Saving %s:%d" % (self._name, index))

    def _make_non_partial(self, new_path):
        """
        Uso interno. Actualiza los campos cuando un archivo deja de ser parcial
        """
        assert self._partial
        self._partial = False
        self._path = new_path

        
# Si se corre este archivo directamente en vez de importar, corremos tests
if __name__ == "__main__":
    import doctest
    doctest.testmod()
    
